/** @file res_shader.cpp
    @brief Shaders */

#include "res/res_shader.h"

using namespace G3;

//==============================================================================
// Constructors
//==============================================================================
G3::SHADER_BASE::SHADER_BASE () {
    Log = &G3::MainLog;
    programID = 0;
    shaderFormat = SF_UNKNOWN;
    valid = false;

    numImageUnits = 0;
}

G3::VERTEX_SHADER::VERTEX_SHADER () {
}

G3::FRAGMENT_SHADER::FRAGMENT_SHADER () {
}

//==============================================================================
// Destructors
//==============================================================================
G3::VERTEX_SHADER::~VERTEX_SHADER () {
    Clear ();
}

G3::FRAGMENT_SHADER::~FRAGMENT_SHADER () {
    Clear ();
}

//==============================================================================
// Gets the file format from its extension
//==============================================================================
G3::SHADER_FORMAT G3::SHADER_BASE::GetFileFormat (const std::string &Filename) {
    if (shaderFormat != SF_UNKNOWN && filename == Filename)
        return shaderFormat;

    filename = Filename;

    TOKENLIST tkl = G3::Tokenize (Filename, ".");

    if (tkl.Tokens [tkl.size () - 1].compare ("asm") == 0) {
        shaderFormat = SF_ASM;
        return SF_ASM;
    }

    if (tkl.Tokens [tkl.size () - 1].compare ("cg") == 0 ||
        tkl.Tokens [tkl.size () - 1].compare ("arb") == 0) {
        shaderFormat = SF_ASM_ASCII;
        return SF_ASM_ASCII;
    }

    if (tkl.Tokens [tkl.size () - 1].compare ("glsl") == 0 ||
        tkl.Tokens [tkl.size () - 1].compare ("vert") == 0 ||
        tkl.Tokens [tkl.size () - 1].compare ("frag") == 0) {
        shaderFormat = SF_GLSL;
        return SF_GLSL;
    }
}

//==============================================================================
// Clears the shader
//==============================================================================
void G3::SHADER_BASE::Clear () {
    if (valid)
        glDeleteProgramsARB (1, &programID);

    samplerImageUnits.clear ();
    valid = false;
}

//==============================================================================
// Loads a vertex shader file
//==============================================================================
bool G3::VERTEX_SHADER::Load (const std::string &Filename) {
    int error = 0;
    valid = false;

    std::string fPath = MainFilesys.GetPath (Filename.c_str ());

    STREAM *Stream = MainFileHandler.ReadWholeFile (fPath.c_str ());

    if (!Stream) {
        if (Log)
            Log->Report ("VERTEX_SHADER::Load: The file \"%s\" doesn't exist..\n", fPath.c_str ());
        return false;
    }

    if (!IsSupported (Filename)) {
        switch (shaderFormat) {
            case SF_ASM:
            case SF_ASM_ASCII:
                if (Log)
                    Log->Report ("WARNING! ARB_vertex_program is not supported. Ignoring \"%s\"\n", Filename.c_str ());
                break;
            case SF_GLSL:
                if (Log)
                    Log->Report ("WARNING! ARB_vertex_shader is not supported. Ignoring \"%s\"\n", Filename.c_str ());
                break;
            case SF_UNKNOWN:
                if (Log)
                    Log->Report ("WARNING! Unknown shader format. Ignoring \"%s\"\n", Filename.c_str ());
                break;
        }
        return false;
    }

    if (shaderFormat == SF_ASM || shaderFormat == SF_ASM_ASCII) {
        if (Log)
            Log->Report ("VERTEX_SHADER::Load: Uploading \"%s\" as an ARB_vertex_program..\n", Filename.c_str ());

        glGenProgramsARB (1, &programID);
        glBindProgramARB (GL_VERTEX_PROGRAM_ARB, programID);

        byte *data = NULL;
        ulong ssize = Stream->Tell (STREAM::STREAM_END);
        try {
            data = new byte [ssize + 1];
        } catch (...) {
            if (Log)
                Log->Report ("VERTEX_SHADER::Load: Failed to allocate memory for a data array of %d bytes..\n", Stream->Tell (STREAM::STREAM_END) + 1);
            return false;
        }

        long size = Stream->Read (data, ssize);

        // Stream::Read doesn't add any 0 bytes, because it's used for binary aswell
        data [ssize] = 0;

#ifdef G3_SHADER_DEBUG
        if (Log)
            Log->Report ("VERTEX_SHADER::Load: ASM Shader (%d bytes) \"%s\"\n", ssize, data);
#endif

        if (shaderFormat == SF_ASM)
            glProgramStringARB (GL_VERTEX_PROGRAM_ARB, GL_PROGRAM_FORMAT_ARB, size, (char *) data);
        else
            glProgramStringARB (GL_VERTEX_PROGRAM_ARB, GL_PROGRAM_FORMAT_ASCII_ARB, size, (char *) data);

        try {
            delete [] data;
        } catch (...) {
            if (Log)
                Log->Report ("VERTEX_SHADER::Load: Failed to delete the data array..\n");
            return false;
        }

        error = glGetError ();

        if (error != GL_NO_ERROR) {
            if (Log)
                Log->Report ("VERTEX_SHADER: The shader \"%s\" threw error %d\n", Filename.c_str (), error);
            return false;
        }

        valid = true;
    } else if (shaderFormat == SF_GLSL) {
        if (Log)
            Log->Report ("VERTEX_SHADER::Load: Uploading \"%s\" as an ARB_vertex_shader..\n", Filename.c_str ());

        const char *shaderStrings [1];
        char str [4096];

        programID = glCreateShaderObjectARB (GL_VERTEX_SHADER_ARB);

        byte *data = NULL;
        ulong ssize = Stream->Tell (STREAM::STREAM_END);
        try {
            data = new byte [ssize + 1];
        } catch (...) {
            if (Log)
                Log->Report ("VERTEX_SHADER::Load: Failed to allocate memory for a data array of %d bytes..\n", Stream->Tell (STREAM::STREAM_END) + 1);
            return false;
        }

        long size = Stream->Read (data, ssize);

        // Stream::Read doesn't add any 0 bytes, because it's used for binary aswell
        data [ssize] = 0;

#ifdef G3_SHADER_DEBUG
        if (Log)
            Log->Report ("VERTEX_SHADER::Load: GLSL Shader (%d bytes) \"%s\"\n", ssize, data);
#endif

        shaderStrings [0] = (char *) data;

        glShaderSourceARB (programID, 1, shaderStrings, NULL);
        glCompileShaderARB (programID);

        try {
            delete [] data;
        } catch (...) {
            if (Log)
                Log->Report ("VERTEX_SHADER::Load: Failed to delete the data array..\n");
            return false;
        }

        glGetObjectParameterivARB (programID, GL_OBJECT_COMPILE_STATUS_ARB, &error);

        if (!error) {
            glGetInfoLogARB (programID, sizeof (str), NULL, str);

            if (Log)
                Log->Report ("VERTEX_SHADER: The shader \"%s\" failed to compile:\n%s", Filename.c_str (), str);
            return false;
        }

        if (Log)
            Log->Report ("VERTEX_SHADER::Load: Uploading completed successfully\n");
    }

    return true;
}

//==============================================================================
// Applies the vertex shader
//==============================================================================
void G3::VERTEX_SHADER::Apply () {
    if (shaderFormat == SF_ASM || shaderFormat == SF_ASM_ASCII) {
        glEnable (GL_VERTEX_PROGRAM_ARB);
        glBindProgramARB (GL_VERTEX_PROGRAM_ARB, programID);
    }
}

//==============================================================================
// Unapplies the vertex shader
//==============================================================================
void G3::VERTEX_SHADER::UnApply () {
    if (shaderFormat == SF_ASM || shaderFormat == SF_ASM_ASCII) {
        glDisable (GL_VERTEX_PROGRAM_ARB);
    }
}

//==============================================================================
// Checks if the needed vertex shader extensions are supported
//==============================================================================
bool G3::VERTEX_SHADER::IsSupported (const std::string &Filename) {
    switch (GetFileFormat (Filename)) {
        case SF_ASM:
        case SF_ASM_ASCII:
            return GLEE_ARB_vertex_program;

        case SF_GLSL:
            return GLEE_ARB_vertex_shader;

        default: break;
    }
    return false;
}

//==============================================================================
// Loads a fragment shader
//==============================================================================
bool G3::FRAGMENT_SHADER::Load (const std::string &Filename) {
    int error = 0;
    valid = false;

    std::string fPath = G3::MainFilesys.GetPath (Filename.c_str ());

    STREAM *Stream = G3::MainFileHandler.ReadWholeFile (fPath.c_str ());

    if (!Stream) {
        if (Log)
            Log->Report ("FRAGMENT_SHADER::Load: The file \"%s\" doesn't exist..\n", fPath.c_str ());
        return false;
    }

    if (!IsSupported (Filename)) {
        switch (shaderFormat) {
            case SF_ASM:
            case SF_ASM_ASCII:
                if (Log)
                    Log->Report ("WARNING! ARB_fragment_program is not supported. Ignoring \"%s\"\n", Filename.c_str ());
                break;
            case SF_GLSL:
                if (Log)
                    Log->Report ("WARNING! ARB_fragment_shader is not supported. Ignoring \"%s\"\n", Filename.c_str ());
                break;
            case SF_UNKNOWN:
                if (Log)
                    Log->Report ("WARNING! Unknown shader format. Ignoring \"%s\"\n", Filename.c_str ());
                break;
        }
        return false;
    }

    if (shaderFormat == SF_ASM || shaderFormat == SF_ASM_ASCII) {
        if (Log)
            Log->Report ("FRAGMENT_SHADER::Load: Uploading \"%s\" as an ARB_fragment_program..\n", Filename.c_str ());

        glGenProgramsARB (1, &programID);
        glBindProgramARB (GL_FRAGMENT_PROGRAM_ARB, programID);

        byte *data = NULL;
        ulong ssize = Stream->Tell (STREAM::STREAM_END);
        try {
            data = new byte [ssize + 1];
        } catch (...) {
            if (Log)
                Log->Report ("FRAGMENT_SHADER::Load: Failed to allocate memory for a data array of %d bytes..\n", Stream->Tell (STREAM::STREAM_END) + 1);
            return false;
        }

        long size = Stream->Read (data, ssize);

        // Stream::Read doesn't add any 0 bytes, because it's used for binary aswell
        data [ssize] = 0;

#ifdef G3_SHADER_DEBUG
        if (Log)
            Log->Report ("FRAGMENT_SHADER::Load: ASM Shader (%d bytes) \"%s\"\n", ssize, data);
#endif

        if (shaderFormat == SF_ASM)
            glProgramStringARB (GL_FRAGMENT_PROGRAM_ARB, GL_PROGRAM_FORMAT_ARB, size, (char *) data);
        else
            glProgramStringARB (GL_FRAGMENT_PROGRAM_ARB, GL_PROGRAM_FORMAT_ASCII_ARB, size, (char *) data);

        try {
            delete [] data;
        } catch (...) {
            if (Log)
                Log->Report ("FRAGMENT_SHADER::Load: Failed to free a data array..\n");
            return false;
        }


        error = glGetError ();

        if (error != GL_NO_ERROR) {
            if (Log)
                Log->Report ("FRAGMENT_SHADER: The shader \"%s\" threw error %d\n", Filename.c_str (), error);
            return false;
        }

        valid = true;
    } else if (shaderFormat == SF_GLSL) {
        if (Log)
            Log->Report ("FRAGMENT_SHADER::Load: Uploading \"%s\" as an ARB_fragment_shader..\n", Filename.c_str ());

        const char *shaderStrings [1];
        char str [4096];

        programID = glCreateShaderObjectARB (GL_FRAGMENT_SHADER_ARB);

        byte *data = NULL;
        ulong ssize = Stream->Tell (STREAM::STREAM_END);
        try {
            data = new byte [ssize + 1];
        } catch (...) {
            if (Log)
                Log->Report ("FRAGMENT_SHADER::Load: Failed to allocate memory for a data array of %d bytes..\n", Stream->Tell (STREAM::STREAM_END) + 1);
            return false;
        }

        long size = Stream->Read (data, ssize);

        // Stream::Read doesn't add any 0 bytes, because it's used for binary aswell
        data [ssize] = 0;

#ifdef G3_SHADER_DEBUG
        if (Log)
            Log->Report ("FRAGMENT_SHADER::Load: GLSL Shader (%d bytes) \"%s\"\n", ssize, data);
#endif

        shaderStrings [0] = (char *) data;

        glShaderSourceARB (programID, 1, shaderStrings, NULL);
        glCompileShaderARB (programID);

        try {
            delete [] data;
        } catch (...) {
            if (Log)
                Log->Report ("FRAGMENT_SHADER::Load: Failed to free a data array..\n");
            return false;
        }

        glGetObjectParameterivARB (programID, GL_OBJECT_COMPILE_STATUS_ARB, &error);

        if (!error) {
            glGetInfoLogARB (programID, sizeof (str), NULL, str);

            if (Log)
                Log->Report ("FRAGMENT_SHADER: The shader \"%s\" failed to compile:\n%s", Filename.c_str (), str);
            return false;
        }

        if (Log)
            Log->Report ("FRAGMENT_SHADER::Load: Uploading completed successfully\n");
    }

    return true;
}

//==============================================================================
// Applies the fragment shader
//==============================================================================
void G3::FRAGMENT_SHADER::Apply () {
    if (shaderFormat == SF_ASM || shaderFormat == SF_ASM_ASCII) {
        glEnable (GL_FRAGMENT_PROGRAM_ARB);
        glBindProgramARB (GL_FRAGMENT_PROGRAM_ARB, programID);
    }
}

//==============================================================================
// Unapplies the fragment shader
//==============================================================================
void G3::FRAGMENT_SHADER::UnApply () {
    if (shaderFormat == SF_ASM || shaderFormat == SF_ASM_ASCII) {
        glDisable (GL_FRAGMENT_PROGRAM_ARB);
    }
}

//==============================================================================
// Sets an uniform sampler2D
//==============================================================================
bool G3::SHADER_BASE::SetTexture2D (uint program_id, const std::string &name, H_TEXTURE handle) {
    if (name.empty () || handle.IsNull ())
        return false;

    std::pair <std::map<std::string, int>::iterator, bool> nipair;
    int imgUnit = numImageUnits;

    // Try to insert it into the map
    nipair = samplerImageUnits.insert (std::pair<std::string, int> (name, imgUnit));

    // The entry already exists ?
    if (!nipair.second)
        imgUnit = nipair.first->second;
    else
        numImageUnits++;

    // Activate the texture
    glActiveTextureARB (GL_TEXTURE0_ARB + imgUnit);
    // Errorcheck
    int error = glGetError ();
    if (error != GL_NO_ERROR) {
        if (Log)
            Log->Report ("SHADER_BASE::SetTexture2D (glActiveTextureARB): OpenGL threw error %d (%s)..\n", error, gluErrorString (error));
    }

    // Ask for the uniform variable location
    int location = glGetUniformLocationARB (program_id, name.c_str ());
    // Errorcheck
    error = glGetError ();
    if (error != GL_NO_ERROR) {
        if (Log)
            Log->Report ("SHADER_BASE::SetTexture2D (glGetUniformLocationARB): OpenGL threw error %d (%s)..\n", error, gluErrorString (error));
    }

    // Bind the texture
    glEnable (GL_TEXTURE_2D);
    glBindTexture (GL_TEXTURE_2D, TextureManager.GetTextureID (handle));
    // Errorcheck
    error = glGetError ();
    if (error != GL_NO_ERROR) {
        if (Log)
            Log->Report ("SHADER_BASE::SetTexture2D (glBindTexture): OpenGL threw error %d (%s)..\n", error, gluErrorString (error));
    }

    // Set the uniform variable
    glUniform1iARB (location, imgUnit);
    // Errorcheck
    error = glGetError ();
    if (error != GL_NO_ERROR) {
        if (Log)
            Log->Report ("SHADER_BASE::SetTexture2D (glUniform1iARB): OpenGL threw error %d (%s)..\n", error, gluErrorString (error));
    }

    glActiveTextureARB (GL_TEXTURE0_ARB);
    glDisable (GL_TEXTURE_2D);
}

//==============================================================================
// Sets an uniform float
//==============================================================================
bool G3::SHADER_BASE::SetFloat (uint program_id, const std::string &name, float value) {
    if (name.empty ())
        return false;

    // Ask for the uniform variable location
    int location = glGetUniformLocationARB (program_id, name.c_str ());
    // Errorcheck
    int error = glGetError ();
    if (error != GL_NO_ERROR) {
        if (Log)
            Log->Report ("SHADER_BASE::SetFloat (glGetUniformLocationARB): OpenGL threw error %d (%s)..\n", error, gluErrorString (error));
    }

    // Set the uniform variable
    glUniform1fARB (location, value);
    // Errorcheck
    error = glGetError ();
    if (error != GL_NO_ERROR) {
        if (Log)
            Log->Report ("SHADER_BASE::SetFloat (glUniform1fARB): OpenGL threw error %d (%s)..\n", error, gluErrorString (error));
    }
}

//==============================================================================
// Sets an uniform vec3
//==============================================================================
bool G3::SHADER_BASE::SetVec3 (uint program_id, const std::string &name, const VEC3F &value) {
    if (name.empty ())
        return false;

    // Ask for the uniform variable location
    int location = glGetUniformLocationARB (program_id, name.c_str ());
    // Errorcheck
    int error = glGetError ();
    if (error != GL_NO_ERROR) {
        if (Log)
            Log->Report ("SHADER_BASE::SetVec3 (glGetUniformLocationARB): OpenGL threw error %d (%s)..\n", error, gluErrorString (error));
    }

    // Set the uniform variable
    glUniform3fARB (location, value.x, value.y, value.z);
    // Errorcheck
    error = glGetError ();
    if (error != GL_NO_ERROR) {
        if (Log)
            Log->Report ("SHADER_BASE::SetVec3 (glUniform3fARB): OpenGL threw error %d (%s)..\n", error, gluErrorString (error));
    }
}

//==============================================================================
// Sets an uniform vec4
//==============================================================================
bool G3::SHADER_BASE::SetVec4 (uint program_id, const std::string &name, const COL4F &value) {
    if (name.empty ())
        return false;

    // Ask for the uniform variable location
    int location = glGetUniformLocationARB (program_id, name.c_str ());
    // Errorcheck
    int error = glGetError ();
    if (error != GL_NO_ERROR) {
        if (Log)
            Log->Report ("SHADER_BASE::SetVec4 (glGetUniformLocationARB): OpenGL threw error %d (%s)..\n", error, gluErrorString (error));
    }

    // Set the uniform variable
    glUniform4fARB (location, value.r, value.g, value.b, value.a);
    // Errorcheck
    error = glGetError ();
    if (error != GL_NO_ERROR) {
        if (Log)
            Log->Report ("SHADER_BASE::SetVec4 (glUniform4fARB): OpenGL threw error %d (%s)..\n", error, gluErrorString (error));
    }
}

//==============================================================================
// Checks if the needed fragment shader extensions are supported
//==============================================================================
bool G3::FRAGMENT_SHADER::IsSupported (const std::string &Filename) {
    switch (GetFileFormat (Filename)) {
        case SF_ASM:
        case SF_ASM_ASCII:
            return GLEE_ARB_fragment_program;

        case SF_GLSL:
            return GLEE_ARB_fragment_shader;

        default: break;
    }
    return false;
}
